"""
Preprocess dataset for countdown task - given a target number and N numbers, generate equations to reach target
"""

import re
import os
from datasets import Dataset, load_dataset
from random import randint, seed, choice
from typing import List, Tuple
from tqdm import tqdm
from verl.utils.hdfs_io import copy, makedirs
import argparse
import pandas as pd
import json


prompt_template = """ A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.
User: {question}
Assistant: <think>
"""

data_source = 'agent_cdm'

train_dir = "/llm/nankai/xuyang_space/data/AM-DeepSeek-Distilled-40M/agent_cdm/raw/data20250514_extract/train.jsonl"
test_dir = "/llm/nankai/xuyang_space/data/AM-DeepSeek-Distilled-40M/agent_cdm/raw/data20250601_test/all.jsonl"

save_dir = "/llm/nankai/xuyang_space/data/AM-DeepSeek-Distilled-40M/agent_cdm/processed/data20250603"



train_dataset = load_dataset('json', data_files=train_dir)
test_dataset = load_dataset('json', data_files=test_dir)



train_dataset = train_dataset["train"]
test_dataset = test_dataset["train"]


def make_map_fn(split):
    def process_fn(example, idx):
        if "question" not in example:
            models = ["Qwen2.5-7B", "Qwen2.5-7B-Instruct", "qwen_7b_r1"]
            question = example[models[0]]["question"]
            solution = example[models[0]]["raw_input"]["answer"]
        else:
            question = example["question"]
            solution = example["answer"]

        question_content = prompt_template.format(
            question=question
        )
        data = {
            "data_source": data_source,
            "prompt": [{
                "role": "user",
                "content": question_content,
            }],
            "ability": "math",
            "reward_model": {
                "style": "rule",
                "ground_truth": solution
            },
            "extra_info": {
                'split': split,
                'index': idx,
            }
        }
        return data
    return process_fn

train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)

# train_dataset = pd.DataFrame(train_dataset)
# test_dataset = pd.DataFrame(test_dataset)

train_dataset.to_parquet(os.path.join(save_dir, 'train.parquet'))
test_dataset.to_parquet(os.path.join(save_dir, 'test.parquet'))

